"""
编码器-解码器架构的Routing Transformer实现

这个模块实现了基于Routing Transformer的编码器-解码器架构，主要特点：
1. 支持独立配置编码器和解码器参数
2. 自动处理上下文掩码和注意力机制
3. 支持序列生成和条件生成
4. 处理可逆网络的辅助损失
"""

import re
from inspect import isfunction
import torch
from torch import nn
from routing_transformer.routing_transformer import (
    RoutingTransformerLM,
    update_kmeans_on_backwards,
)
from routing_transformer.autoregressive_wrapper import AutoregressiveWrapper

# 编码器和解码器参数前缀
ENC_PREFIX = "enc_"
DEC_PREFIX = "dec_"


def default(x, d):
    """
    默认值处理函数

    参数:
        x: 输入值
        d: 默认值或生成默认值的函数

    返回:
        如果x为None，返回默认值；否则返回x
    """
    if x is None:
        return d if not isfunction(d) else d()
    return x


def group_dict_by_key(cond, d):
    """
    根据条件函数对字典键进行分组

    参数:
        cond: 条件函数，接受键作为参数
        d: 要分组的字典

    返回:
        两个字典的元组，分别包含满足和不满足条件的项
    """
    return_val = [dict(), dict()]
    for key in d.keys():
        match = bool(cond(key))
        ind = int(not match)
        return_val[ind][key] = d[key]
    return (*return_val,)


def string_begins_with(prefix, str):
    """
    检查字符串是否以指定前缀开始

    参数:
        prefix: 前缀字符串
        str: 要检查的字符串
    """
    return bool(re.match(f"^{prefix}", str))


def group_by_key_prefix(prefix, d):
    """
    根据键前缀对字典进行分组

    参数:
        prefix: 前缀字符串
        d: 要分组的字典
    """
    return group_dict_by_key(lambda x: string_begins_with(prefix, x), d)


def group_by_key_prefix_and_remove_prefix(prefix, d):
    """
    根据键前缀分组并移除前缀

    参数:
        prefix: 前缀字符串
        d: 要处理的字典

    返回:
        (移除前缀后的匹配字典, 不匹配的字典)
    """
    kwargs_with_prefix, kwargs = group_dict_by_key(
        lambda x: string_begins_with(prefix, x), d
    )
    kwargs_without_prefix = dict(
        map(lambda x: (x[0][len(prefix) :], x[1]), tuple(kwargs_with_prefix.items()))
    )
    return kwargs_without_prefix, kwargs


def extract_enc_dec_kwargs(kwargs):
    """
    从kwargs中提取编码器和解码器的参数

    参数:
        kwargs: 包含编码器和解码器参数的字典

    返回:
        (编码器参数, 解码器参数, 其他参数)
    """
    enc_kwargs, kwargs = group_by_key_prefix_and_remove_prefix(ENC_PREFIX, kwargs)
    dec_kwargs, kwargs = group_by_key_prefix_and_remove_prefix(DEC_PREFIX, kwargs)
    return enc_kwargs, dec_kwargs, kwargs


def extract_and_set_enc_dec_kwargs(kwargs):
    """
    提取并设置编码器和解码器的参数，处理特殊情况

    参数:
        kwargs: 参数字典

    返回:
        (编码器参数, 解码器参数, 其他参数)
    """
    enc_kwargs, dec_kwargs, kwargs = extract_enc_dec_kwargs(kwargs)
    if "input_mask" in enc_kwargs:
        dec_kwargs.setdefault("context_mask", enc_kwargs["input_mask"])
    return enc_kwargs, dec_kwargs, kwargs


class RoutingTransformerEncDec(nn.Module):
    """
    Routing Transformer的编码器-解码器模型

    实现了基于Routing Transformer的序列到序列转换模型，具有以下特点：
    1. 独立的编码器和解码器配置
    2. 支持因果注意力和上下文注意力
    3. 处理可逆网络的特殊情况
    4. 支持序列生成和条件生成

    参数:
        dim: 模型维度
        ignore_index: 忽略的索引值（用于损失计算）
        pad_value: 填充值
        **kwargs: 其他参数，可包含编码器和解码器的特定参数
    """

    def __init__(self, dim, ignore_index=None, pad_value=0, **kwargs):
        super().__init__()
        ignore_index = default(ignore_index, pad_value)
        enc_kwargs, dec_kwargs, _ = extract_enc_dec_kwargs(kwargs)

        assert (
            "return_embedding" not in enc_kwargs
        ), "you cannot manually set the return embeddings flag for the encoder"
        assert (
            "dim" not in dec_kwargs and "dim" not in enc_kwargs
        ), "you must set the dim for both encoder and decoder"

        # 设置编码器和解码器的基本参数
        enc_kwargs["dim"] = dec_kwargs["dim"] = dim
        enc_kwargs["return_embeddings"] = True
        dec_kwargs["causal"] = True
        dec_kwargs["receives_context"] = True
        enc_kwargs["_register_kmeans_update"] = dec_kwargs[
            "_register_kmeans_update"
        ] = False

        enc_kwargs.setdefault("window_size", 256)
        dec_kwargs.setdefault("window_size", 256)

        # 创建编码器和解码器
        enc = RoutingTransformerLM(**enc_kwargs)
        dec = RoutingTransformerLM(**dec_kwargs)

        self.enc = enc
        self.dec = AutoregressiveWrapper(
            dec, ignore_index=ignore_index, pad_value=pad_value
        )

        # 处理解码器可逆性的特殊情况
        self.dec_reversible = dec_kwargs.pop("reversible", False)

        if self.dec_reversible:
            print(
                "Warning! Due to an issue with reversible nets and encoder auxiliary losses, you must explicitly call backwards on the encoder auxiliary loss, which is supplied as the second element of the returned tuple on forward"
            )

        self._handle = None
        self.register_kmeans_update()

    def cancel_kmeans_update(self):
        """取消k-means更新的注册"""
        if self._handle is None:
            return
        self._handle.remove()
        self._handle = None

    def register_kmeans_update(self):
        """注册k-means更新"""
        self.cancel_kmeans_update()
        return update_kmeans_on_backwards(self)

    @torch.no_grad()
    def generate(self, seq_in, seq_out_start, max_seq_len=None, **kwargs):
        """
        生成序列

        参数:
            seq_in: 输入序列
            seq_out_start: 输出序列的起始部分
            max_seq_len: 最大序列长度
            **kwargs: 其他参数

        返回:
            生成的序列
        """
        max_seq_len = default(max_seq_len, self.dec.max_seq_len)
        enc_kwargs, dec_kwargs, kwargs = extract_and_set_enc_dec_kwargs(kwargs)
        context, _ = self.enc(seq_in, **enc_kwargs)
        return self.dec.generate(
            seq_out_start, max_seq_len, context=context, **{**dec_kwargs, **kwargs}
        )

    def forward(
        self,
        seq_in,
        seq_out,
        return_loss=False,
        randomly_truncate_sequence=False,
        **kwargs,
    ):
        """
        前向传播

        参数:
            seq_in: 输入序列
            seq_out: 目标输出序列
            return_loss: 是否返回损失
            randomly_truncate_sequence: 是否随机截断序列
            **kwargs: 其他参数

        返回:
            (损失, 辅助损失)
        """
        enc_kwargs, dec_kwargs, kwargs = extract_and_set_enc_dec_kwargs(kwargs)
        context, enc_aux_loss = self.enc(seq_in, **enc_kwargs)
        loss = self.dec(
            seq_out,
            return_loss=return_loss,
            randomly_truncate_sequence=randomly_truncate_sequence,
            context=context,
            aux_loss=enc_aux_loss,
            **dec_kwargs,
        )

        # 处理解码器可逆性的特殊情况
        if self.dec_reversible:
            return loss, enc_aux_loss

        aux_loss = torch.tensor(0.0, requires_grad=True)
        loss = loss + enc_aux_loss
        return loss, aux_loss
